﻿/**********************************************************************************************************************/
/* See License.txt file for details.                                                                                  */
/**********************************************************************************************************************/

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Microsoft.Win32;
using System.Runtime.InteropServices;
using System.IO;
using System.Runtime.ConstrainedExecution;
using Microsoft.Win32.SafeHandles;
using System.Security.Permissions;
using System.Security.Principal;

namespace CodePlex.SqlServer.LocalDB.API
{
    public class LocalDBException : Exception
    {
        public LocalDBException(string message)
            : base(message)
        {
        }
    }

    public static class LocalDBAPI
    {
        #region Safe handles and Win32 imports

        /// <summary>
        /// See http://msdn.microsoft.com/msdnmag/issues/05/10/Reliability/ 
        /// for more about safe handles.
        /// </summary>
        [SecurityPermission(SecurityAction.LinkDemand, UnmanagedCode = true)]
        sealed class SafeLibraryHandle : SafeHandleZeroOrMinusOneIsInvalid
        {
            /// <summary>
            /// Create safe library handle
            /// </summary>
            private SafeLibraryHandle() : base(true) { }

            /// <summary>
            /// Release handle
            /// </summary>
            protected override bool ReleaseHandle()
            {
                return NativeMethod.FreeLibrary(handle);
            }
        }

        /// <summary>
        /// Native methods
        /// </summary>
        static class NativeMethod
        {
            [DllImport("kernel32", CharSet = CharSet.Auto, SetLastError = true)]
            public static extern SafeLibraryHandle LoadLibrary(string fileName);

            [ReliabilityContract(Consistency.WillNotCorruptState, Cer.Success)]
            [DllImport("kernel32", SetLastError = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool FreeLibrary(IntPtr hModule);

            [DllImport("kernel32", EntryPoint = "GetProcAddress")]
            public static extern IntPtr GetProcAddress(SafeLibraryHandle hModule, String procname);
        }

        #endregion Safe handles and Win32 imports

        #region Internal pointers and delegates

        // Unmanaged resource. CLR will ensure SafeHandles get freed, without 
        // requiring a finalizer on this class.
        // TODO: confirm this works with a static member as well (all samples are showing non-static members)
        static SafeLibraryHandle DllHandle;

        const string LocalDbApiKeyName = @"SOFTWARE\Microsoft\Microsoft SQL Server Local DB\Installed Versions";
        const string LocalDbDllLocationValueName = @"InstanceAPIPath";

        const int MAX_LOCALDB_INSTANCE_NAME_LENGTH = 128;
        const int MAX_LOCALDB_INSTANCE_NAME_BUFFER_LENGTH = MAX_LOCALDB_INSTANCE_NAME_LENGTH + 1;
        const int LOCALDB_MAX_SQLCONNECTION_BUFFER_SIZE = 260;
        const int MAX_STRING_SID_LENGTH = 186;

        const uint FACILITY_CODE_MASK = 0x3fff0000;
        const uint LOCALDB_FACILITY_CODE = 0x09C50000;

        const uint LOCALDB_TRUNCATE_ERR_MESSAGE = 0x0001;
        const uint S_OK = 0x00000000;
        const uint LOCALDB_ERROR_INVALID_OPERATION = 0x89C50113;
        const uint LOCALDB_ERROR_INSUFFICIENT_BUFFER = 0x89C50114;

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        delegate uint LocalDBGetInstancesDelegate(
            [In, Out, MarshalAs(UnmanagedType.LPArray, ArraySubType=UnmanagedType.U2)] char[] pInstanceNames,
            [In, Out] ref uint lpdwNumberOfInstances);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        delegate uint LocalDBGetInstanceInfoDelegate(
            [MarshalAs(UnmanagedType.LPWStr)] string wszInstanceName,
            [In, Out] ref _LocalDBInstanceInfo pInstanceInfo,
            uint dwInstanceInfoSize);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        delegate uint LocalDBStartInstanceDelegate(
            [MarshalAs(UnmanagedType.LPWStr)] string pInstanceName,
            uint dwFlags,
            [In, Out, MarshalAs(UnmanagedType.LPArray, ArraySubType = UnmanagedType.U2)] char[] wszSqlConnection,
            [In, Out] ref uint lpcchSqlConnection);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        delegate uint LocalDBStopInstanceDelegate(
            [MarshalAs(UnmanagedType.LPWStr)] string pInstanceName,
            uint dwFlags,
            uint ulTimeout);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        delegate uint LocalDBShareInstanceDelegate(
            IntPtr pOwnerSID,
            [MarshalAs(UnmanagedType.LPWStr)] string pInstancePrivateName,
            [MarshalAs(UnmanagedType.LPWStr)] string pInstanceSharedName,
            uint dwFlags);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        delegate uint LocalDBUnshareInstanceDelegate(
            [MarshalAs(UnmanagedType.LPWStr)] string pInstanceSharedName,
            uint dwFlags);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        delegate uint LocalDBFormatMessageDelegate(
            uint hrLocalDBCode,
            uint dwFlags,
            uint dwLanguageId,
            [In, Out, MarshalAs(UnmanagedType.LPArray, ArraySubType = UnmanagedType.U2)] char[] wszMessage,
            [In, Out] ref uint lpcchMessage);
        

        static LocalDBGetInstancesDelegate LocalDBGetInstances = null;
        static LocalDBGetInstanceInfoDelegate LocalDBGetInstanceInfo = null;
        static LocalDBStartInstanceDelegate LocalDBStartInstance = null;
        static LocalDBStopInstanceDelegate LocalDBStopInstance = null;
        static LocalDBShareInstanceDelegate LocalDBShareInstance = null;
        static LocalDBUnshareInstanceDelegate LocalDBUnshareInstance = null;
        static LocalDBFormatMessageDelegate LocalDBFormatMessage = null;

        [StructLayout(LayoutKind.Sequential)]
        struct _LocalDBInstanceInfo
        {
	        uint cbLocalDBInstanceInfoSize;
            [MarshalAs(UnmanagedType.ByValArray, ArraySubType = UnmanagedType.U2, SizeConst = MAX_LOCALDB_INSTANCE_NAME_BUFFER_LENGTH)]
	        char[] wszInstanceName;
            bool bExists;
	        bool bConfigurationCorrupted;
	        bool bIsRunning;
	        uint dwMajor;
        	uint dwMinor;
	        uint dwBuild;
	        uint dwRevision;
            System.Runtime.InteropServices.ComTypes.FILETIME ftLastStartDateUTC; //Per: http://go.microsoft.com/fwlink/?linkid=14202
            [MarshalAs(UnmanagedType.ByValArray, ArraySubType = UnmanagedType.U2, SizeConst = LOCALDB_MAX_SQLCONNECTION_BUFFER_SIZE)]
	        char[] wszConnection;
	        bool bIsShared;
            [MarshalAs(UnmanagedType.ByValArray, ArraySubType = UnmanagedType.U2, SizeConst = MAX_LOCALDB_INSTANCE_NAME_BUFFER_LENGTH)]
	        char[] wszSharedInstanceName;
            [MarshalAs(UnmanagedType.ByValArray, ArraySubType = UnmanagedType.U2, SizeConst = MAX_STRING_SID_LENGTH + 1)]
            char[] wszOwnerSID;
	        bool bIsAutomatic;

            internal void Init()
            {
                cbLocalDBInstanceInfoSize = (uint)Marshal.SizeOf(typeof(_LocalDBInstanceInfo));
                wszInstanceName = new char[MAX_LOCALDB_INSTANCE_NAME_BUFFER_LENGTH];
                wszConnection = new char[LOCALDB_MAX_SQLCONNECTION_BUFFER_SIZE];
                wszSharedInstanceName = new char[MAX_LOCALDB_INSTANCE_NAME_BUFFER_LENGTH];
                wszOwnerSID = new char[MAX_STRING_SID_LENGTH + 1];
            }

            internal LocalDBInstanceInfo ToPublicStruct()
            {
                return new LocalDBInstanceInfo(
                    new string(wszInstanceName).TrimEnd('\0'),
                    new string(wszSharedInstanceName).TrimEnd('\0'),
                    bExists,
                    bIsRunning,
                    bIsShared,
                    new string(wszConnection).TrimEnd('\0'),
                    new string(wszOwnerSID).TrimEnd('\0') // TODO: change string into some SID structure
                    );
            }
        };

        static string GetDllPath()
        {
            var localDbApiKey = Registry.LocalMachine.OpenSubKey(LocalDbApiKeyName);
            if (localDbApiKey == null)
                throw new LocalDBException("SQL Server Express LocalDB is not installed.");

            // TODO: extend to cover future versions

            var localDbApiKey11 = localDbApiKey.OpenSubKey("11.0");
            if (localDbApiKey11 == null)
                throw new LocalDBException("SQL Server Express 2012 LocalDB is not installed and this is the only version of SQL Server Express LocalDB supported so far by LocalDB Companion.");

            string localDbDllLocation = localDbApiKey11.GetValue(LocalDbDllLocationValueName) as string;
            if (localDbApiKey11 == null)
                throw new LocalDBException("SQL Server Express 2012 LocalDB installation is corrupt. Please repair it before using LocalDB Companion.");

            return localDbDllLocation;
        }

        static void InitializeDll()
        {
            if (DllHandle == null)
            {
                DllHandle = NativeMethod.LoadLibrary(GetDllPath());
                if (DllHandle.IsInvalid)
                {
                    int hr = Marshal.GetHRForLastWin32Error();
                    Marshal.ThrowExceptionForHR(hr);
                }
            }
        }

        static void InitLocalDBGetInstances()
        {
            if (LocalDBGetInstances == null)
            {
                LocalDBGetInstances = GetUnmanagedFunction<LocalDBGetInstancesDelegate>("LocalDBGetInstances");
            }
        }

        static void InitLocalDBGetInstanceInfo()
        {
            if (LocalDBGetInstanceInfo == null)
            {
                LocalDBGetInstanceInfo = GetUnmanagedFunction<LocalDBGetInstanceInfoDelegate>("LocalDBGetInstanceInfo");
            }
        }

        static void InitLocalDBStartInstance()
        {
            if (LocalDBStartInstance == null)
            {
                LocalDBStartInstance = GetUnmanagedFunction<LocalDBStartInstanceDelegate>("LocalDBStartInstance");
            }
        }

        static void InitLocalDBStopInstance()
        {
            if (LocalDBStopInstance == null)
            {
                LocalDBStopInstance = GetUnmanagedFunction<LocalDBStopInstanceDelegate>("LocalDBStopInstance");
            }
        }

        static void InitLocalDBShareInstance()
        {
            if (LocalDBShareInstance == null)
            {
                LocalDBShareInstance = GetUnmanagedFunction<LocalDBShareInstanceDelegate>("LocalDBShareInstance");
            }
        }

        static void InitLocalDBUnshareInstance()
        {
            if (LocalDBUnshareInstance == null)
            {
                LocalDBUnshareInstance = GetUnmanagedFunction<LocalDBUnshareInstanceDelegate>("LocalDBUnshareInstance");
            }
        }

        static void InitLocalDBFormatMessage()
        {
            if (LocalDBFormatMessage == null)
            {
                LocalDBFormatMessage = GetUnmanagedFunction<LocalDBFormatMessageDelegate>("LocalDBFormatMessage");
            }
        }

        /// <summary>
        /// Dynamically lookup a function in the dll via kernel32!GetProcAddress.
        /// </summary>
        /// <param name="functionName">
        /// raw name of the function in the export table.
        /// </param>
        /// <returns>
        /// null if function is not found. Else a delegate to the unmanaged 
        /// function.
        /// </returns>
        /// <remarks>
        /// GetProcAddress results are valid as long as the dll is not yet 
        /// unloaded. This is very very dangerous to use since you need to 
        /// ensure that the dll is not unloaded until after you're done with any 
        /// objects implemented by the dll. For example, if you get a delegate 
        /// that then gets an IUnknown implemented by this dll, you can not 
        /// dispose this library until that IUnknown is collected. Else, you may 
        /// free the library and then the CLR may call release on that IUnknown 
        /// and it will crash.
        /// </remarks>
        static TDelegate GetUnmanagedFunction<TDelegate>(string functionName)
            where TDelegate : class
        {
            InitializeDll();

            IntPtr p = NativeMethod.GetProcAddress(DllHandle, functionName);

            // Failure is a common case, especially for adaptive code.
            if (p == IntPtr.Zero)
            {
                return null;
            }

            Delegate function = Marshal.GetDelegateForFunctionPointer(p, typeof(TDelegate));

            // Ideally, we'd just make the constraint on TDelegate be
            // System.Delegate, but compiler error CS0702 
            // (constrained can't be System.Delegate)
            // prevents that. So we make the constraint system.object and do the
            // cast from object-->TDelegate.
            object o = function;

            return (TDelegate)o;
        }

        static void FreeDll()
        {
            // TODO: implement and hook to something :-)
        }

        public static string ThrowExceptionForErrorCode(uint errorCode)
        {
            InitLocalDBFormatMessage();

            // Check if the error is a general Windows HRESULT or LocalDB specific one
            if ((errorCode & FACILITY_CODE_MASK) != LOCALDB_FACILITY_CODE)
            {
                // For Windows HRESULTS use typed exceptions
                Marshal.ThrowExceptionForHR((int)errorCode);

                // For regular Windows error codes just throw a general error message
                throw new Exception(string.Format("General error, code {0}.", errorCode));
            }

            // For LocalDB errors use LocalDBFormatMessage to get the error message
            uint dwBufferLength = 2048;
            char[] wszBuffer = new char[dwBufferLength];
            uint hrResult = LocalDBFormatMessage(errorCode, LOCALDB_TRUNCATE_ERR_MESSAGE, 0, wszBuffer, ref dwBufferLength);

            if (hrResult != S_OK)
            {
                throw new Exception(string.Format("Unexpected error in LocalDB API, code {0:X8}.", hrResult));
            }

            throw new LocalDBException(new string(wszBuffer, 0, (int)dwBufferLength).TrimEnd('\0'));
        }

        #endregion Internal pointers and delegates

        #region Public LocalDB API Wrappers

        public static string[] GetInstances()
        {
            InitLocalDBGetInstances();

            uint dwInstanceNamesCount = 0;

            uint hrResult = LocalDBGetInstances(null, ref dwInstanceNamesCount);
            if (hrResult != LOCALDB_ERROR_INSUFFICIENT_BUFFER)
            {
                ThrowExceptionForErrorCode(hrResult);
            }

            char[] instanceNamesBuffer = new char[MAX_LOCALDB_INSTANCE_NAME_BUFFER_LENGTH * dwInstanceNamesCount];

            hrResult = LocalDBGetInstances(instanceNamesBuffer, ref dwInstanceNamesCount);

            if (hrResult != S_OK)
            {
                ThrowExceptionForErrorCode(hrResult);
            }

            var results = new string[dwInstanceNamesCount];

            for (uint i = 0; i < dwInstanceNamesCount; i++)
            {
                results[i] = new string(instanceNamesBuffer, (int)(i * MAX_LOCALDB_INSTANCE_NAME_BUFFER_LENGTH), (int)MAX_LOCALDB_INSTANCE_NAME_LENGTH).TrimEnd('\0');
            }

            return results;
        }

        public static LocalDBInstanceInfo GetInstanceInfo(string instanceName)
        {
            InitLocalDBGetInstanceInfo();

            _LocalDBInstanceInfo instanceInfoInt = new _LocalDBInstanceInfo();
            instanceInfoInt.Init();

            uint hrResult = LocalDBGetInstanceInfo(instanceName, ref instanceInfoInt, (uint)Marshal.SizeOf(typeof(_LocalDBInstanceInfo)));

            if (hrResult != S_OK)
            {
                ThrowExceptionForErrorCode(hrResult);
            }

            return instanceInfoInt.ToPublicStruct();
        }

        public static void StartInstance(string instanceName)
        {
            InitLocalDBStartInstance();

            uint dwConnectionAddressLength = LOCALDB_MAX_SQLCONNECTION_BUFFER_SIZE;
            char[] connectionAdressChars = new char[dwConnectionAddressLength];
            uint hrResult = LocalDBStartInstance(instanceName, 0, connectionAdressChars, ref dwConnectionAddressLength);

            if (hrResult != S_OK)
            {
                ThrowExceptionForErrorCode(hrResult);
            }
        }

        public static void StopInstance(string instanceName, uint timeoutMiliseconds)
        {
            InitLocalDBStopInstance();

            uint hrResult = LocalDBStopInstance(instanceName, 0, timeoutMiliseconds);

            if (hrResult != S_OK)
            {
                ThrowExceptionForErrorCode(hrResult);
            }
        }

        public static void ShareInstance(string privateInstanceName, string publicInstanceName)
        {
            InitLocalDBShareInstance();

            uint hrResult = LocalDBShareInstance(IntPtr.Zero, privateInstanceName, publicInstanceName, 0);

            if (hrResult != S_OK)
            {
                ThrowExceptionForErrorCode(hrResult);
            }
        }

        public static void UnshareInstance(string publicInstanceName)
        {
            InitLocalDBUnshareInstance();

            uint hrResult = LocalDBUnshareInstance(publicInstanceName, 0);

            if (hrResult != S_OK)
            {
                ThrowExceptionForErrorCode(hrResult);
            }
        }

        #endregion Public LocalDB API Wrappers
    }
}
